"""
DWT距离计算模块
使用离散小波变换计算两个信号的距离
"""
import numpy as np
import pywt
from typing import List, Tuple, Dict, Optional
import warnings

class DWTCalculator:
    def __init__(self, wavelet: str = 'db4', levels: int = 3):
        """
        初始化DWT计算器
        
        Args:
            wavelet: 小波基函数类型
            levels: 分解层数
        """
        self.wavelet = wavelet
        self.levels = levels
        
        # 验证小波是否有效
        if wavelet not in pywt.wavelist():
            raise ValueError(f"不支持的小波类型: {wavelet}")
    
    def _pad_signal(self, signal: np.ndarray, target_length: int) -> np.ndarray:
        """
        填充信号到目标长度
        
        Args:
            signal: 输入信号
            target_length: 目标长度
            
        Returns:
            填充后的信号
        """
        if len(signal) >= target_length:
            return signal[:target_length]
        
        # 使用对称填充
        pad_length = target_length - len(signal)
        pad_left = pad_length // 2
        pad_right = pad_length - pad_left
        
        return np.pad(signal, (pad_left, pad_right), mode='symmetric')
    
    def _align_signals(self, signal1: np.ndarray, signal2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        对齐两个信号长度
        
        Args:
            signal1: 信号1
            signal2: 信号2
            
        Returns:
            对齐后的信号对
        """
        max_length = max(len(signal1), len(signal2))
        
        # 确保长度是2的幂次（对DWT更友好）
        power_of_2 = 2 ** int(np.ceil(np.log2(max_length)))
        
        signal1_padded = self._pad_signal(signal1, power_of_2)
        signal2_padded = self._pad_signal(signal2, power_of_2)
        
        return signal1_padded, signal2_padded
    
    def dwt_decompose(self, signal: np.ndarray) -> List[np.ndarray]:
        """
        对信号进行小波分解
        
        Args:
            signal: 输入信号
            
        Returns:
            小波系数列表 [cA_n, cD_n, cD_n-1, ..., cD_1]
        """
        try:
            coeffs = pywt.wavedec(signal, self.wavelet, level=self.levels)
            return coeffs
        except Exception as e:
            print(f"小波分解失败: {e}")
            return [signal]  # 返回原信号作为单个系数
    
    def calculate_coeffs_distance(self, coeffs1: List[np.ndarray], coeffs2: List[np.ndarray]) -> float:
        """
        计算小波系数之间的欧氏距离
        
        Args:
            coeffs1: 信号1的小波系数
            coeffs2: 信号2的小波系数
            
        Returns:
            系数间的总距离
        """
        if len(coeffs1) != len(coeffs2):
            raise ValueError("小波系数列表长度不匹配")
        
        total_distance = 0.0
        
        for c1, c2 in zip(coeffs1, coeffs2):
            # 确保系数长度相同
            min_len = min(len(c1), len(c2))
            c1_trim = c1[:min_len]
            c2_trim = c2[:min_len]
            
            # 计算欧氏距离
            distance = np.sqrt(np.sum((c1_trim - c2_trim) ** 2))
            total_distance += distance
        
        return total_distance
    
    def calculate_weighted_distance(self, coeffs1: List[np.ndarray], coeffs2: List[np.ndarray], 
                                  weights: Optional[List[float]] = None) -> float:
        """
        计算加权小波系数距离
        
        Args:
            coeffs1: 信号1的小波系数
            coeffs2: 信号2的小波系数
            weights: 各层权重，如果为None则使用默认权重
            
        Returns:
            加权距离
        """
        if len(coeffs1) != len(coeffs2):
            raise ValueError("小波系数列表长度不匹配")
        
        num_levels = len(coeffs1)
        
        # 默认权重：低频部分权重更大
        if weights is None:
            weights = [2.0] + [1.0] * (num_levels - 1)  # 近似系数权重更大
        
        if len(weights) != num_levels:
            raise ValueError("权重数量与系数层数不匹配")
        
        weighted_distance = 0.0
        
        for i, (c1, c2) in enumerate(zip(coeffs1, coeffs2)):
            # 确保系数长度相同
            min_len = min(len(c1), len(c2))
            c1_trim = c1[:min_len]
            c2_trim = c2[:min_len]
            
            # 计算欧氏距离并加权
            distance = np.sqrt(np.sum((c1_trim - c2_trim) ** 2))
            weighted_distance += weights[i] * distance
        
        return weighted_distance
    
    def dwt_distance(self, signal1: np.ndarray, signal2: np.ndarray, 
                     use_weights: bool = True, weights: Optional[List[float]] = None) -> Dict:
        """
        计算两个信号的DWT距离
        
        Args:
            signal1: 信号1（如标准曲线的采样点）
            signal2: 信号2（如用户录制的采样点）
            use_weights: 是否使用加权距离
            weights: 自定义权重
            
        Returns:
            包含距离信息的字典
        """
        try:
            # 对齐信号长度
            signal1_aligned, signal2_aligned = self._align_signals(signal1, signal2)
            
            # 小波分解
            coeffs1 = self.dwt_decompose(signal1_aligned)
            coeffs2 = self.dwt_decompose(signal2_aligned)
            
            # 计算距离
            if use_weights:
                distance = self.calculate_weighted_distance(coeffs1, coeffs2, weights)
                distance_type = "weighted"
            else:
                distance = self.calculate_coeffs_distance(coeffs1, coeffs2)
                distance_type = "unweighted"
            
            return {
                'distance': distance,
                'distance_type': distance_type,
                'wavelet': self.wavelet,
                'levels': self.levels,
                'signal1_length': len(signal1),
                'signal2_length': len(signal2),
                'aligned_length': len(signal1_aligned),
                'coeffs_count': len(coeffs1),
                'success': True
            }
            
        except Exception as e:
            return {
                'distance': float('inf'),
                'distance_type': 'error',
                'error': str(e),
                'success': False
            }
    
    def calculate_similarity_score(self, distance: float, max_distance: float = None) -> float:
        """
        将距离转换为相似度分数 (0-100)
        
        Args:
            distance: DWT距离
            max_distance: 最大距离（用于归一化），如果为None则使用经验值
            
        Returns:
            相似度分数 (0-100)
        """
        if max_distance is None:
            # 经验最大距离值，可根据实际数据调整
            max_distance = 1000.0
        
        if distance >= max_distance:
            return 0.0
        
        # 距离越小，相似度越高
        similarity = (1 - distance / max_distance) * 100
        return max(0.0, min(100.0, similarity))
    
    def batch_compare(self, standard_signal: np.ndarray, 
                     user_signals: List[np.ndarray]) -> List[Dict]:
        """
        批量计算标准信号与多个用户信号的距离
        
        Args:
            standard_signal: 标准信号
            user_signals: 用户信号列表
            
        Returns:
            距离结果列表
        """
        results = []
        
        for i, user_signal in enumerate(user_signals):
            result = self.dwt_distance(standard_signal, user_signal)
            result['user_signal_index'] = i
            
            if result['success']:
                result['similarity_score'] = self.calculate_similarity_score(result['distance'])
            else:
                result['similarity_score'] = 0.0
            
            results.append(result)
        
        return results
    
    def get_available_wavelets(self) -> Dict[str, List[str]]:
        """
        获取可用的小波类型
        
        Returns:
            按类别分组的小波列表
        """
        return {
            'Daubechies': [w for w in pywt.wavelist() if w.startswith('db')],
            'Biorthogonal': [w for w in pywt.wavelist() if w.startswith('bior')],
            'Coiflets': [w for w in pywt.wavelist() if w.startswith('coif')],
            'Discrete Meyer': [w for w in pywt.wavelist() if w.startswith('dmey')],
            'Haar': ['haar'],
            'Others': [w for w in pywt.wavelist() if not any(w.startswith(p) for p in ['db', 'bior', 'coif', 'dmey']) and w != 'haar']
        }

# 使用示例
if __name__ == "__main__":
    calculator = DWTCalculator()
    
    # 模拟测试数据：标准信号（一声上升调）
    t = np.linspace(0, 1, 20)
    standard_signal = 200 + 50 * t
    
    # 模拟用户信号（与标准信号相似但有偏差）
    user_signal1 = 205 + 45 * t + np.random.normal(0, 5, 20)  # 相似度较高
    user_signal2 = 180 + 70 * t + np.random.normal(0, 10, 20) # 相似度较低
    
    print("DWT距离计算测试:")
    print("=" * 50)
    
    # 测试单个比较
    result1 = calculator.dwt_distance(standard_signal, user_signal1)
    print(f"用户信号1 vs 标准信号:")
    print(f"  距离: {result1['distance']:.3f}")
    print(f"  相似度: {calculator.calculate_similarity_score(result1['distance']):.1f}%")
    print(f"  小波: {result1['wavelet']}, 层数: {result1['levels']}")
    
    result2 = calculator.dwt_distance(standard_signal, user_signal2)
    print(f"\n用户信号2 vs 标准信号:")
    print(f"  距离: {result2['distance']:.3f}")
    print(f"  相似度: {calculator.calculate_similarity_score(result2['distance']):.1f}%")
    
    # 测试批量比较
    print(f"\n批量比较结果:")
    batch_results = calculator.batch_compare(standard_signal, [user_signal1, user_signal2])
    for i, result in enumerate(batch_results):
        print(f"  信号{i+1}: 距离={result['distance']:.3f}, 相似度={result['similarity_score']:.1f}%")
    
    # 显示可用小波
    print(f"\n推荐小波类型:")
    wavelets = calculator.get_available_wavelets()
    print(f"  Daubechies: {wavelets['Daubechies'][:5]}...")  # 只显示前5个
